"""
股票价格预测模块

该模块实现股票价格预测功能，主要包含：
1. 从数据库加载最新股票数据
2. 数据预处理
3. 加载训练好的模型
4. 生成预测结果
"""
import torch
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from config import DB_CONFIG
from model_trainer import StockTransformer

def predict_stock_price(date=None):
    """
    股票价格预测函数
    
    功能:
        1. 加载训练好的模型
        2. 从数据库获取最新股票数据
        3. 进行数据预处理
        4. 生成预测结果
        
    返回:
        float: 预测的股票价格
    """
    try:
        # 初始化数据库连接
        engine = create_engine(
            f"mysql+mysqlconnector://{DB_CONFIG['user']}:{DB_CONFIG['password']}@"
            f"{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}"
        )
        
        # 加载指定日期前6天的股票数据(包含当天)
        if not date:
            date = pd.read_sql("SELECT MAX(trade_date) FROM stock_index", engine).iloc[0,0]
            
        df = pd.read_sql(
            """
            SELECT * FROM stock_index 
            WHERE trade_date <= %s
            ORDER BY trade_date DESC 
            LIMIT 6
            """, 
            engine,
            params=(date,)
        )
        print(f'df: {df}')
        print(f'从数据库加载到{len(df)}条数据')
        # 数据预处理(与训练时一致)
        # 提取与训练时相同的42维特征
        features = df[['open', 'high', 'low', 'close', 'pre_close', 'vol', 'amount']].values
        features = np.concatenate([features[i] for i in range(len(features))])  # 展平为42维数组
        print(f'输入特征: {features}')
        
        # 加载模型
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model = StockTransformer().to(device)
        model.load_state_dict(torch.load('best_model.pth'))
        model.eval()
        
        # 生成预测
        with torch.no_grad():
            inputs = torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(device)  # 添加batch维度
            inputs = inputs.unsqueeze(0)  # 添加序列维度
            prediction = model(inputs).item()
        
        # 反标准化预测结果
        last_close = df.iloc[-1]['close']
        prediction = prediction * df['close'].std() + df['close'].mean()
        
        print(f'预测结果: {prediction:.2f}')
        return prediction
        
    except Exception as e:
        print(f'预测失败: {str(e)}')
        raise


if __name__ == "__main__":
    # 示例：预测指定日期的股票价格
    predict_stock_price('2025-03-24')